package com.cxk.ml

import org.apache.spark.ml.evaluation.RegressionEvaluator
import org.apache.spark.ml.recommendation.ALS
import org.apache.spark.sql.types.{IntegerType, StructField, StructType}
import org.apache.spark.sql.{DataFrame, Row, SparkSession}

object ALSRecommend {

  def als(spark: SparkSession, train: DataFrame, user_index: DataFrame, item_index: DataFrame): Unit = {
    val data = train.select("msno", "song_id", "target")

    val df = data.join(user_index, Seq("msno"), "left")
      .join(item_index, Seq("song_id"), "left")
      .drop("msno", "song_id")

    //Spark的ALS数据需要user,item,ratting，且user,item必须整数
    val Array(training, test) = spark.createDataFrame(df.rdd
      .map(row => Row(row.getString(0).toInt, row.getString(1).toInt, row.getString(2).toInt)),
      StructType(df.schema.map(s => StructField(s.name, IntegerType))))
      .randomSplit(Array(0.8, 0.2))

    //使用训练数据训练模型
    //这里的ALS是import org.apache.spark.ml.recommendation.ALS，不是mllib中的
    //setMaxiter设置最大迭代次数
    //setRegParam设置正则化参数
    //setUserCol设置用户id列名
    //setItemCol设置物品列名
    //setRatingCol设置打分列名
    val als = new ALS()

      //    als.setRank(10)
      //      .setMaxIter(5)
      //      .setRegParam(0.01)
      .setUserCol("user_id")
      .setItemCol("item_id")
      .setRatingCol("target")

    //fit给输出的数据，训练模型，fit返回的是ALSModel类
    val model = als.fit(training)

    //使用测试数据计算模型的误差平方和
    //transform方法把数据dataset换成dataframe类型，预测数据
    val predictions = model.transform(test)
    println("predicted.count==============================================" + predictions.count())
    predictions.show()

    //RegressionEvaluator这个类是用户评估预测效果的，预测值与原始值
    //这个setLabelCol要和als设置的setRatingCol一致，不然会报错
    //RegressionEvaluator的setPredictionCol必须是prediction因为，ALSModel的默认predictionCol也是prediction
    //如果要修改的话必须把ALSModel和RegressionEvaluator一起修改
    //model.setPredictionCol("prediction")和evaluator.setPredictionCol("prediction")
    //setMetricName这个方法，评估方法的名字，一共有哪些呢？
    //rmse-平均误差平方和开根号
    //mse-平均误差平方和
    //mae-平均距离（绝对）
    //r2-没用过不知道
    //这里建议就是用rmse就好了，其他的基本都没用，当然还是要看应用场景，这里是预测分值就是用rmse。如果是预测距离什么的mae就不从，看场景哈
    val evaluator = new RegressionEvaluator()
      .setMetricName("rmse")
      .setLabelCol("target")
      .setPredictionCol("prediction")

    val rmse = evaluator.evaluate(predictions)
    println("Root-mean-square error = " + rmse)

    model.recommendForAllItems(test.select("user_id").first().getInt(0)).show()
  }

}
